#!/usr/bin/env python
import socket
import itertools
import sys

from dm_reg import *
from dm_utils import *
from openocd import *

def init_progbuf(ocd):
    asm = [
            0x000104b7, # lui s1, 0x10
            0x0404849b, # addiw s1, s1, 0x40
            0x7b149073, # csrw dpc, s1
            0x7b002473, # csrr s0, dcsr
            0x00346413, # ori s0, s0 3 # change to m-mode
            0x7b041073, # csrw dcsr, s0
            0x00100073  # ebreak
            ]
    for i, data in enumerate(asm):
        ocd.send("riscv dmi_write {} {}".format(0x20 + i, data))

def halt_hart(ocd, hart):
    while True:
        ocd.dmi_write_fields(dmcontrol, hartsello = hart, resumereq = 0, haltreq = 1)
        ocd.dmi_read_fields(dmstatus, v=False)
        if dmstatus.allhalted == 1:
            break
        print("not halt, retry")
    ocd.dmi_write_fields(dmcontrol, haltreq = 0)

def resume_hart(ocd, hart):
    while True:
        ocd.dmi_write_fields(dmcontrol, hartsello = hart, resumereq = 1, haltreq = 0)
        ocd.dmi_read_fields(dmstatus, v=False)
        if dmstatus.allrunning == 1:
            break
        print("not resume, retry")
    ocd.dmi_write_fields(dmcontrol, resumereq = 0)


def restart_hart(ocd, hart):
    # Access Register
    ocd.send("riscv dmi_write 0x17 0x361001") # regno = sp, postexec, transfer, 64-bit
    abstractcs_ = ocd.send("riscv dmi_read 0x16")
    abstractcs.bits = int(abstractcs_, 16)
    if abstractcs_ != "0x10000002":
        print("Exec error: abstracts " + abstractcs_)
        print_fields(abstractcs)


if __name__ == "__main__":
    harts = sys.argv[1:]
    with OpenOcd() as ocd:
        # Progbuf
        init_progbuf(ocd)
        for hart in harts:
            print("start hart " + hart)
            hart = int(hart)
            halt_hart(ocd, hart)
            restart_hart(ocd, hart)
            resume_hart(ocd, hart)
